import re


class Meta(object):
    def __init__(self, code, clsdt, Wclsprc, Wsmvosd, Wretwd):
        # 股票交易代码
        self.code = code

        #
        self.date = clsdt
        self.value = Wclsprc
        self.total_value = Wsmvosd
        self.rate = Wretwd

    def __str__(self):
        return "code = {},date = {},value = {},total_value = {},rate = {}".format(self.code, self.date, self.value,
                                                                                  self.total_value, self.rate)


class Stock(object):
    def __init__(self, code, metas=None):
        self.code = code
        if metas is None:
            self.metas = []
        else:
            self.metas = metas

    def append(self, meta):
        self.metas.append(meta)

    # def append(self, metas):
    #     for me in metas:
    #         self.metas.append(me)

    def clean_data(self):
        """
        清洗数据，剔除新上市骨片前八周的数据
        :param data_list:按时间顺序排列的某一code对应的全部数据
        :return:
        """
        if len(self.metas) == 0:
            print("股票{}没有数据".format(self.code))

        first = self.metas[0]
        # 如果不是从第一周开始剔除前八周,否则保留全部数据，
        if first.date != "2006/1/6":
            self.metas = [] if len(self.metas) <= 8 else self.metas[8:]

    def get_pre_rate_by_date(self, date, weeks, weekMap, indexMap):
        """
        给定日期，和周数，过去该时间范围内的累计收益率
        :param data:
        :param weeks:
        :param weekMap:
        :param indexMap:
        :return:
        """

        rate = 1
        # 开始下标
        end = weekMap[date]
        # 开始下标
        start = end - weeks

        for data in self.metas:
            if start <= week_map[data.date] < end:
                rate *= (data.rate + 1)

        return rate - 1

    def get_next_rate_by_date(self, date, weeks, weekMap, indexMap):
        """
        给定日期，和周数，过去该时间范围内的累计收益率
        :param data:
        :param weeks:
        :param weekMap:
        :param indexMap:
        :return:
        """

        rate = 1
        # 开始下标
        start = weekMap[date]
        # 开始下标
        end = start + weeks - 1

        for data in self.metas:
            if start <= week_map[data.date] < end:
                rate *= (data.rate + 1)

        return rate - 1

    def is_appear(self, data, weeks, week_map, index_map):
        """
        判断该股票是否出现在指定时间范围内
        :param data:
        :param weeks:
        :param week_map:
        :param index_map:
        :return:
        """
        # 开始下标
        end = week_map[data]
        # 开始下标
        start = end - weeks

        for data in self.metas:
            if start <= week_map[data.date] < end:
                return True

        return False


def getWeekMap(data_list):
    """
    获取两个map,前者可以给定日期，得到他是第几周，后者可以给定第几周得到日期
    :param data_list:
    :return:
    """
    # 存放日期到下标的映射
    weekMap = {}
    # 存放下标到日期的映射
    indexMap = {}
    # 存放日期列表
    timeList = []

    count = 1
    lastdate = ""
    for item in data_list:
        if item.date != lastdate:
            timeList.append(item.date)
            weekMap[item.date] = count
            indexMap[count] = item.date
            count += 1
            lastdate = item.date

    return weekMap, indexMap, timeList


def get_group(stock_list, date, weeks, weekMap, indexMap, rate):
    # 先从stock_list过滤出在 data前weeks到date直接出现的股票
    stock_list = list(filter(lambda x: x.is_appear(date, weeks, weekMap, indexMap), stock_list))

    size = len(stock_list)
    n = int(size * rate)

    rate_list = []
    # 获取过去weeks的累计收益率，并排序

    for stock in stock_list:
        rate_list.append(
            (stock.code, stock.get_pre_rate_by_date(date, weeks, weekMap=weekMap, indexMap=indexMap), stock))

    # 按收益率排序
    rate_list = list(map(lambda x: x[2], sorted(rate_list, key=lambda x: x[1], reverse=True)))

    # 获取前百分之10和后百分之10的数据
    return rate_list[:n], rate_list[-1 * n:]


def get_avg_rate(rate_list):
    size = len(rate_list)
    return sum(rate_list) / size


if __name__ == '__main__':

    import sys

    parm = sys.argv
    j, k, rate = int(parm[1]), int(parm[2]), float(parm[3])
    print("观察期：{}周，持有期：{}周，百分比：{}%".format(j, k, rate * 100))
    print("计算中...")

    source = open("newdata.csv", encoding='utf-8')

    # 1.数据预处理，获取每个stock的所有交易数据

    data_list = []

    for line in source.readlines():
        line = line.replace("\n", "").split(",")

        try:
            line = Meta(re.findall("[0-9]+", line[0])[0], line[1], float(line[2]), float(line[3]), float(line[4]))
        except Exception as e:
            line = Meta(re.findall("[0-9]+", line[0])[0], line[1], float(line[2]), float(line[3]), 0.0)

        data_list.append(line)

    week_map, index_map, timeList = getWeekMap(data_list)

    stock_map = {}

    for item in data_list:
        stock = stock_map.get(item.code, Stock(item.code))
        stock.append(item)
        stock_map[item.code] = stock

    # 清理每个stock的交易数据，如果是早已上市的股票，不清理，如果新上市搞得股票，清理前8周的数据
    for stock_key in stock_map:
        stock_map[stock_key].clean_data()

    stock_list = list(map(lambda x: x[1], stock_map.items()))

    R_list = []
    RW_list = []
    RL_list = []
    # 每次循环选定一个t，计算对于给定（j,k）下，对于持有期的股票的收益率
    for t in timeList:
        # 获得在 [t-j,t)之间排名前rate和后rate的股票
        R0, R1 = get_group(stock_list, t, j, week_map, index_map, rate)

        if len(R0) == 0 or len(R1) == 0:
            continue

        # 计算赢家组合中每个股票在[t,t+k)的累计收益率，并计算所有股票累计收益率的平均值
        R0_rate_list = []
        for R0_item in R0:
            R0_rate_list.append(R0_item.get_next_rate_by_date(t, k + 1, week_map, index_map))
        RW = sum(R0_rate_list) / len(R0_rate_list)

        # 计算输家组合中每个股票在[t,t+k)的累计收益率，并计算所有股票累计收益率的平均值
        R1_rate_list = []
        for R1_item in R1:
            R1_rate_list.append(R1_item.get_next_rate_by_date(t, k + 1, week_map, index_map))
        RL = sum(R1_rate_list) / len(R1_rate_list)

        RW_RL = RW - RL

        R_list.append((t, RW, RL, RW_RL))
        RW_list.append(RW)
        RL_list.append(RL)

    print(RW_list)
    print(RL_list)
    RW_avg = sum(RW_list) / len(RW_list)
    RL_avg = sum(RL_list) / len(RL_list)

    RW_RL_avg = RW_avg - RL_avg

    print(RW_avg - RL_avg)
